import pathlib

import matplotlib.pyplot
import torch

from checkpoint import Checkpoint
from checkpoint_manager import CheckpointManager
from scaler import Scaler
from simple_pendulum_dataset import SimplePendulumDataset
from spinn import Spinn


def main():
    problems = [
        (1, "not_damped", Spinn()),
        (2, "under_damping", Spinn()),
        (3, "over_damping", Spinn(30, 1))
    ]

    for subplot_index, title, model in problems:
        matplotlib.pyplot.subplot(1, len(problems), subplot_index)
        model_path = pathlib.Path(f"./checkpoints/{title}")
        data_path = pathlib.Path(f"./data/{title}.pt")

        dataset = SimplePendulumDataset.load(data_path)
        t_min = float(dataset.t_min())
        t_max = float(dataset.t_max())

        _, checkpoint = CheckpointManager(model_path).last_existed_loaded(Checkpoint.load)
        assert checkpoint is not None
        model.scalar = Scaler(t_min, t_max)
        model.load_state_dict(checkpoint.model_state())
        model.eval()

        matplotlib.pyplot.scatter(dataset.t(), dataset.theta(), label="Data")
        t = torch.linspace(t_min, t_max, 1000)
        matplotlib.pyplot.plot(t, torch.detach(model(t)), label="Spinn")
        matplotlib.pyplot.legend(loc="upper right")
        matplotlib.pyplot.title(title)
        matplotlib.pyplot.xlabel(r"$t$")
        matplotlib.pyplot.ylabel(r"$\theta$")

    matplotlib.pyplot.subplots_adjust(left=0.07, right=0.955, wspace=0.25)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
